package com.pai4j.connect.websocket.handler;

import com.pai4j.common.enums.RoleCodeEnum;
import com.pai4j.common.enums.messagequeue.MessageBroadChannelEnum;
import com.pai4j.common.util.JsonUtil;
import com.pai4j.connect.service.AccountService;
import com.pai4j.connect.service.AsyncMessageService;
import com.pai4j.domain.vo.request.mess.ChatMessageRequestVO;
import com.pai4j.domain.vo.response.UserBaseResponseInfoVO;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 *  这里使用的是springboot3.x 的websocket实现,springboot2.x的websocket实现，使用的是注解形式
 *
 * @author: CYM-pai
 * @date: 2025/05/13 17:07
 **/
@Slf4j
public class ChatWebSocketEndPointServletHandler extends TextWebSocketHandler {

    @Autowired
    private AccountService accountService;

    @Autowired
    private AsyncMessageService asyncMessageService;

    /**
     * Map<KEY1, Map<KEY2, WebSocketSession>>
     *
     * KEY1：用于区分场景，如：PC-APP、 PC-CHAT、WAP-APP ，详见：WebSocketBizCodeEnum
     * KEY2: 用户
     * WebSocketSession: 不同场景下用户的WebSocket Session对象 -- 全双工、有状态
     */
    public static Map<String, Map<String, WebSocketSession>> sessions = new ConcurrentHashMap<>();

    /**
     * 连接建立
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 获取路径参数
        Map<String, Object> attributes = session.getAttributes();
        String bizCode = (String) attributes.get("bizCode");
        String userId = (String) attributes.get("userId");
        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (userSessions == null) {
            userSessions = new ConcurrentHashMap<>();
            sessions.put(bizCode, userSessions);
        }

        // 如果同一bizCode下该用户已有连接，关闭旧连接，防止重复连接堆积
        WebSocketSession old = userSessions.get(userId);
        if (old != null && old.isOpen() && !StringUtils.equals(old.getId(), session.getId())) {
            try {
                log.warn("Duplicate WebSocket connection detected, closing old session. userId:{}, oldSessionId:{}, newSessionId:{}", userId, old.getId(), session.getId());
                old.close(CloseStatus.NORMAL);
            } catch (Exception e) {
                log.warn("Failed to close old WebSocket session. userId:{}, oldSessionId:{}", userId, old.getId(), e);
            }
        }

        userSessions.put(userId, session);
        log.info("Chat WebSocket connection ====> userId:{}, sessionId:{}", userId, session.getId());
    }

    /**
     * 收到消息
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        //从WebSocket消息中获取负载（payload），即客户端(前端)发送的原始数据，通常是JSON格式的消息内容。
        String payload = message.getPayload();
        log.debug("Chat WebSocket received message: {}", payload);
        if ("ping".equals(payload)) {
            // 心跳检测
            return;
        }
        //将客户端发送的消息内容解析为ChatMessageRequestVO对象
        ChatMessageRequestVO chatMessage = JsonUtil.fromJson(payload, ChatMessageRequestVO.class);
        if (chatMessage == null) {
            return;
        }
        // 获取对话消息接收人
        String receiverId = chatMessage.getReceiverId();
        UserBaseResponseInfoVO accountInfo = accountService.accountInfo(receiverId);
        if (accountInfo == null) {
            return;
        }
        
        // 获取用户ID和会话ID用于异步处理
        String userId = (String) session.getAttributes().get("userId");
        String sessionId = chatMessage.getSessionId();
        
        // 使用异步处理避免阻塞WebSocket连接
        if (RoleCodeEnum.AIBOT.getCode().equals(accountInfo.getRoleCode())) {
            /**
             * 我的ai助手对话: AGENT对话，走AGENT工作流处理
             */
            asyncMessageService.broadcastMessageWithCallback(
                MessageBroadChannelEnum.QUEUE_CHAT_AI_BOT, 
                chatMessage, 
                userId, 
                sessionId,
                null, // 成功回调可以为空
                () -> sendErrorToClient(session, "AI助手消息发送失败，请稍后重试")
            );
        } else if (RoleCodeEnum.ROBOT.getCode().equals(accountInfo.getRoleCode())) {
            /**
             * 机器人AGENT对话，走AGENT工作流处理
             */
            asyncMessageService.broadcastMessageWithCallback(
                MessageBroadChannelEnum.QUEUE_CHAT_ROBOT, 
                chatMessage, 
                userId, 
                sessionId,
                null, // 成功回调可以为空
                () -> sendErrorToClient(session, "机器人消息发送失败，请稍后重试")
            );
        } else {
            /**
             * 对话消息，走MQ广播到所有消费者（底层支持RocketMQ、Redis等）
             */
            asyncMessageService.broadcastMessageWithCallback(
                MessageBroadChannelEnum.QUEUE_CHAT_MESSAGE, 
                chatMessage, 
                userId, 
                sessionId,
                null, // 成功回调可以为空
                () -> sendErrorToClient(session, "消息发送失败，请稍后重试")
            );
        }
    }

    /**
     * 断开连接
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Map<String, Object> attributes = session.getAttributes();
        String bizCode = (String) attributes.get("bizCode");
        String userId = (String) attributes.get("userId");
        if (StringUtils.isNotBlank(bizCode) && StringUtils.isNotBlank(userId)) {
            Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
            if (MapUtils.isNotEmpty(userSessions)) {
                WebSocketSession existing = userSessions.get(userId);
                if (existing != null && StringUtils.equals(existing.getId(), session.getId())) {
                    userSessions.remove(userId);
                }
                if (userSessions.isEmpty()) {
                    sessions.remove(bizCode);
                }
            }
        }
        log.info("Chat WebSocket connection closed. userId:{}, sessionId:{}, status:{}", userId, session.getId(), status);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Map<String, Object> attributes = session.getAttributes();
        String userId = (String) attributes.get("userId");
        log.warn("Chat WebSocket transport error. userId:{}, sessionId:{}", userId, session.getId(), exception);
        try {
            if (session.isOpen()) {
                session.close(CloseStatus.SERVER_ERROR);
            }
        } catch (Exception e) {
            log.debug("Error while closing session after transport error. sessionId:{}", session.getId(), e);
        }
    }
    /**
     * 单个用户的WebSocket Session对象获取：一对一聊天场景
     *
     * @author: CYM-pai
     * @date: 2025/05/13 17:23
     **/
    public static WebSocketSession getSession(String bizCode, String userId) {
        if (StringUtils.isBlank(bizCode) || StringUtils.isBlank(userId)) {
            return null;
        }
        //当前用户使用的场景下的所有用户WebSocket Session对象
        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return null;
        }
        //获取对应用户id的WebSocket Session对象
        return userSessions.get(userId);
    }
    
    /**
     * 批量获取用户WebSocket Session对象： 多对多群聊场景
     *
     * @author: CYM-pai
     * @date: 2025/05/13 17:24
     **/
    
    public static Map<String, WebSocketSession> getSessions(String bizCode, Set<String> userIds) {
        if (StringUtils.isBlank(bizCode) || CollectionUtils.isEmpty(userIds)) {
            return Collections.emptyMap();
        }
        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return Collections.emptyMap();
        }
        Map<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();
        userIds.forEach(uid -> {
            WebSocketSession ws = userSessions.get(uid);
            if (ws != null) {
                userSessionMap.put(uid, ws);
            }
        });
        return userSessionMap;
    }

    /**
     * 向客户端发送错误消息
     *
     * @param session WebSocket会话
     * @param errorMessage 错误消息
     */
    private void sendErrorToClient(WebSocketSession session, String errorMessage) {
        if (session != null && session.isOpen()) {
            try {
                String errorJson = String.format("{\"type\":\"error\",\"message\":\"%s\"}", errorMessage);
                session.sendMessage(new TextMessage(errorJson));
            } catch (Exception e) {
                log.debug("Failed to send error message to client. sessionId:{}, error:{}", session.getId(), e.getMessage());
            }
        }
    }
}
